from torch.utils.data import Dataset
import torch
import cv2
import os
from typing import Optional,List,Tuple,Dict
import torchvision
import numpy as np
from random import shuffle

CLASS_NAMES:Dict[str,int]={
    '01_TUMOR':0,
    '02_STROMA':1,
    '03_COMPLEX':2,
    '04_LYMPHO':3,
    '05_DEBRIS':4,
    '06_MUCOSA':5,
    '07_ADIPOSE':6,
    '08_EMPTY':7,
}

def getAllDataPaths(root_folder:str)->Dict[str,List[str]]:
    """
        root_folder必须为绝对路径
    """
    assert(os.path.isabs(root_folder))
    allData:Dict[str,List[str]]={}
    # initialize
    for name in CLASS_NAMES:
        allData[name]=[]

    for folder in os.listdir(root_folder):
        if os.path.isdir(os.path.join(root_folder,folder))==False \
            or (folder not in CLASS_NAMES):
            continue
        # root/folder/tif_file
        for tif_file in os.listdir(os.path.join(root_folder,folder)):
            if os.path.isfile(os.path.join(root_folder,folder,tif_file))==False\
                or (tif_file.endswith('.tiff') or tif_file.endswith('.tif'))==False:
                continue
            allData[folder].append(os.path.join(root_folder,folder,tif_file))
    return allData

def getRandomDataPaths(root_folder:str):
    allData=getAllDataPaths(root_folder)
    for key in allData:
        shuffle(allData[key])
    return allData

class cancer5000Dataset(Dataset):
    def __init__(self,allData:Dict[str,List[str]],\
        raw_transform:Optional[torchvision.transforms.transforms.Compose]=None) -> None:
        super().__init__()
        self.all_pairs:Tuple[Tuple[str]]=[]
        for name, listt in allData.items():
            index:int=CLASS_NAMES[name]
            for tif_file in listt:
                self.all_pairs.append((tif_file,index,))
        self.all_pairs=tuple(self.all_pairs)

        self.raw_transform=raw_transform
        if self.raw_transform==None:
            self.raw_transform=torchvision.transforms.ToTensor()
    
    def __getitem__(self, index: int) -> torch.Tensor:
        raw_img:np.ndarray=cv2.imread(self.all_pairs[index][0],cv2.IMREAD_UNCHANGED)
        label:int=self.all_pairs[index][1]
        return self.raw_transform(raw_img), torch.tensor(label)

    def __len__(self):
        return len(self.all_pairs)

def getTrainTestDataset(root_folder:str,\
    raw_transform:Optional[torchvision.transforms.transforms.Compose]=None,\
    testingPercentage:float=0.2)->Tuple[cancer5000Dataset]:
    allData=getRandomDataPaths(root_folder)
    assert(testingPercentage>=0 and testingPercentage <1)
    trainingPercentage:float=1-testingPercentage
    trainData, testData={},{}
    for name,listt in allData.items():
        total=len(listt)
        trainData[name]=listt[:int(total*trainingPercentage)]
        testData[name]=listt[int(total*trainingPercentage):]
    
    return cancer5000Dataset(trainData,raw_transform),\
        cancer5000Dataset(testData,raw_transform),

if __name__ == '__main__':
    print(type(os.path.join('abc','xya')))
    allData=getAllDataPaths('E:/JohnsonProj/Kather_texture_2016_image_tiles_5000')
    print(allData['05_DEBRIS'][2])
    

    dataset=cancer5000Dataset(allData)
    img_arr:torch.Tensor
    img_arr,i=dataset[4500]
    _,i2=dataset[4501]
    print(len(dataset))
    print(img_arr.shape,i)
    print('下一张:',i2)

    trainSet, testSet=\
        getTrainTestDataset('E:/JohnsonProj/Kather_texture_2016_image_tiles_5000')
    print('划分数据集:',len(trainSet),len(testSet))